# -*-coding:utf-8 -*-
import os
import process_file as pf
import tree_sitter_related as ts


current_id = 0



folder_path = '../dataset/test/'
prefix_original = '../../crates_crawler/crates/'
prefix_new = '../../crates_crawler/crates/'



if __name__ == '__main__':
    for filename in os.listdir(folder_path):
        file_path = os.path.join(folder_path, filename)
        if os.path.isfile(file_path) and filename.endswith('.json'):
            crate_path = f'{prefix_new}{filename[:-5]}'
            list_dict_info = pf.read_json(file_path)
            
            for dict_info in list_dict_info:
                FUT_path = pf.replace_prefix(dict_info["FUT_path"], prefix_original, prefix_new)
                rust_code = pf.read_rust_code(FUT_path)
                root_node = ts.get_root_node(rust_code)
                node_FUT = ts.function_str_to_node(root_node, dict_info["FUT_full"])
                if not node_FUT:
                    continue

                list_node_total_function = []
                list_node_test_function = []
                dict_PATH = {}
                crate_name = pf.get_last_directory_name(crate_path)



                # 遍历整个项目的rs文件，先得到dict_unique_node_function
                for root, dirs, files in os.walk(crate_path):
                    for rsfilename in files:
                        rsfile_path = pf.path_normalized(os.path.join(root, rsfilename))
                        if os.path.isfile(rsfile_path) and rsfilename.endswith('.rs'):
                            rust_code = pf.read_rust_code(rsfile_path)
                            root_node = ts.get_root_node(rust_code)

                            # 找到所有#[test]
                            list_node_test_function_temp = ts.get_list_node_test_function(root_node)

                            # 找到所有function_item
                            list_node_total_function_temp = ts.get_list_node_total_function(root_node)
                            for node in list_node_total_function_temp:
                                dict_PATH[id(node)] = rsfile_path
                            list_node_total_function.extend(list_node_total_function_temp)

                            if list_node_test_function_temp:
                                for node in list_node_test_function_temp:
                                    dict_PATH[id(node)] = rsfile_path
                                list_node_test_function.extend(list_node_test_function_temp)
                
                
                
                # 去重list_node_total_function,(函数名，参数数量)都相同则去掉
                dict_unique_node_function = ts.get_dict_unique_node_function_item(list_node_total_function)
                
                # 删除dict_unique_node_function中的测试函数
                if not list_node_test_function:
                    continue
                for node_test_function in list_node_test_function:
                    if not node_test_function or node_test_function.type != 'function_item':
                        continue
                    name_test_function = ts.get_function_name(node_test_function)
                    cnt_para_num = ts.get_function_para_count(node_test_function)
                    if (name_test_function, cnt_para_num) in dict_unique_node_function:
                        del dict_unique_node_function[(name_test_function, cnt_para_num)]

                # 找info
                dict_func_call_ex = ts.get_dict_func_call_ex(node_FUT, dict_unique_node_function)
                dict_call_ex_func_target = ts.get_dict_call_ex_func_target(node_FUT, dict_unique_node_function)
                node_parent = ts.get_node_function_parent(node_FUT)

                # 保存
                list_CoT_FUT_call_ex = []
                list_CoT_FUT_be_call_ex = []
                list_CoT_FUT_parent = []

                if dict_func_call_ex:
                    for key, value in dict_func_call_ex.items():
                        list_CoT_FUT_call_ex.append(value.text.decode('utf-8'))
                if dict_call_ex_func_target:
                    for key, value in dict_call_ex_func_target.items():
                        list_CoT_FUT_be_call_ex.append(value.text.decode('utf-8'))
                if node_parent:
                    list_CoT_FUT_parent.append(node_parent.text.decode('utf-8'))
                dict_info["CoT_FUT_call_ex"] = list_CoT_FUT_call_ex
                dict_info["CoT_FUT_be_call_ex"] = list_CoT_FUT_be_call_ex
                dict_info["CoT_FUT_parent"] = list_CoT_FUT_parent

            
            # 保存回原文件
            pf.save_data_to_json(list_dict_info, file_path)
            current_id += 1
            print(f"{current_id} {filename} success!!!")




